# ./query/get_sql_answer.py
# usage：
# python ./src/query/get_sql_answer.py --query_text "请问2024年各受理渠道工单所占的百分比分别是多少？" --debug
# 请问2023年3月交通运输专题工单反馈的工单量排名前4的事项有哪些？  唯一TOP01
# 请问在2025年的农林牧渔专题中，D区办结工单量相较于A区多出多少（若A区更多则请用负数表示）？  TOP03
# rerank multi-result: python ./src/query/get_sql_answer.py --query_text "请问2024年区县单位平均处理工单时长是多少？"
# example hint: python ./src/query/get_sql_answer.py --debug --query_text "请问2024年公卫医疗专题各区县服务满意率分别是多少？" 


import os,sys
import argparse
project_root = os.path.abspath(os.path.join(os.path.dirname(__file__), '..'))
sys.path.insert(0, project_root)
import time
# 项目内模块导入（相对路径或绝对路径根据你的结构调整）
from model.ai_loader import client
from config import project_config
from query.query_ddl import query_ddl
from query.query_table import query_table
from query.query_sql import query_sql
from query.build_extra_hint import build_extra_hint
from utils import create_sql_template  
from utils import debug_print
from query.query_example import query_example

def get_sql_answer(question: str, llm="low", debug=None) -> str:
    # 这个问题是否有对应的问题模板
    table_name = query_table(question) 
    table_define = None
    example_hint = None
    if table_name is not None:
        table_define = query_sql(table_name)
        example_hint = query_example(question)
        debug_print(debug, f"[BINGO]{table_name}")
    # 这是一个没有见过的新问题
    if table_define is None and llm=="high":  # This is a new question 
        table_define = query_ddl(question)
    if table_define is None :
        return None

    extra_hint = build_extra_hint(question)  # 额外提示信息
    
    sql_template = create_sql_template()
    prompt = sql_template.format_messages(q=question, table_define=table_define, extra_hint=extra_hint, example_hint=example_hint)

    debug_print(debug, f"问题：{question}")
    debug_print(debug, "-" * 40)
    debug_print(debug, f"prompt: {prompt[0].content}")
    debug_print(debug, "-" * 40)

    # success = False
    # resp = None

    max_retries = 3
    attempts = 0
    success = False
    llm_model = project_config.llm_model_low
    if (llm=="high") : llm_model = project_config.llm_model_high
    if (llm=="normal") : llm_model = project_config.llm_model_normal
    debug_print(debug, f"llm_model={llm_model}")
    while not success and attempts < max_retries:
        try:
            resp = client.chat.completions.create(
                model=llm_model , #project_config.llm_model,
                messages=[{"role": "user", "content": prompt[0].content}],
                temperature=0.7,
                top_p=0.5,
                extra_body={"enable_thinking": False, 'stream': False}
            )
            success = True  # 如果调用成功，标记为成功
        except Exception as e:
            attempts += 1
            print(f"调用失败: {e}，剩余重试次数: {max_retries - attempts}")
            if attempts < max_retries:
                time.sleep(3)
            else:
                print("已达到最大重试次数，放弃重试。\n")
                return None
    sql_result = resp.choices[0].message.content
    debug_print(debug, f"SQL：{sql_result}")
    debug_print(debug, "=" * 50)

    return sql_result

def main():
    # 创建参数解析器
    parser = argparse.ArgumentParser(description="提示词大模型生成靠谱的SQL")
    parser.add_argument('--query_text', type=str, required=True, help='需要检索的文本字符串。'    )
    parser.add_argument('--debug', action='store_true', help='启用调试模式，显示详细日志。')
    parser.add_argument('--llm', type=str, default="low", choices=['low', 'normal', 'high'], help='指定大模型级别：low, normal, 或 high。')
    args = parser.parse_args()
    text_to_query = args.query_text
    
    # 将 silent 参数传递给函数
    sql_result = get_sql_answer(text_to_query, args.llm, args.debug)
    print(f"SQL：{sql_result}")
# --- 程序入口 ---
if __name__ == '__main__':
    main()

